require 'thread'

# This class implements a pool of threads. The pool hands out blocks to be
# processed to any available worker thread. If no worker thread is available,
# the pool blocks until any of the workers have become available. The pool
# size can be specified when constructing the pool. For exmaple:
#
#   pool = ThreadPool.new(10)
#   10.times {pool.process {sleep 3}}
#   10.times {pool.process {sleep 3}} # this will block for about 3 seconds.
#   pool.join # this will block for about 3 more seconds.
class ThreadPool
  # The Worker class starts a thread for performing work. A worker thread
  # sleeps until it receives a block for processing, at which time it is
  # considered busy. Once the block has been processed, the worker is
  # considered available again.
  class Worker
    # Starts the worker thread.
    def initialize
      @mutex = Mutex.new
      @thread = Thread.new do
        while true
          sleep 0.001
          block = get_block
          if block
            block.call
            reset_block
          end
        end
      end
    end
    
    # Returns the block to be processed, if any.
    def get_block
      @mutex.synchronize {@block}
    end
    
    # Sets the block to be processed.
    def set_block(block)
      @mutex.synchronize {@block = block}
    end
    
    # Resets the block instance variable to signify that the worker
    # is available.
    def reset_block
      @mutex.synchronize {@block = nil}
    end
    
    # Returns true if the worker is busy processing a block.
    def busy?
      @mutex.synchronize {!@block.nil?}
    end
  end
  
  attr_accessor :max_size
  attr_reader :workers

  # Creates a new pool with a default maximum size of 10.
  def initialize(max_size = 10)
    @max_size = max_size
    @workers = []
    @mutex = Mutex.new
  end
  
  # Returns the number of workers in the pool.
  def size
    @mutex.synchronize {@workers.size}
  end
  
  # Returns true if any of the workers in the pool are busy.
  def busy?
    @mutex.synchronize {@workers.any? {|w| w.busy?}}
  end
  
  # Blocks until all workers have become available.
  def join
    sleep 0.01 while busy?
  end
  
  # Hands the block to an available worker. If no workers are available, this
  # method will block until one becomes available.
  def process(&block)
    wait_for_worker.set_block(block)
  end
  
  # Waits for a worker to become available.
  def wait_for_worker
    while true
      worker = find_available_worker
      return worker if worker
      sleep 0.01
    end
  end
  
  # Returns an available worker or a new worker if the maximum pool size
  # has not been reached.
  def find_available_worker
    @mutex.synchronize {free_worker || create_worker}
  end
  
  # Returns the first available worker.
  def free_worker
    @workers.each {|w| return w unless w.busy?}; nil
  end
  
  # Creates a worker if the maximum pool size has not been reached.
  def create_worker
    return nil if @workers.size >= @max_size
    worker = Worker.new
    @workers << worker
    worker
  end
end
